import pandas as pd
from gmssl import sm2, func
import base64
import binascii
import os
import argparse
import secrets

def setup_keys(use_fixed_key=True):
    """设置SM2密钥对，每次都生成新密钥"""
    # 开发环境示例
    private_key_file = "d:\\source pro\\secret\\sm2_private.key"
    public_key_file = "d:\\source pro\\secret\\sm2_public.key"
    
    # 生成新密钥
    private_key = secrets.token_hex(32)
    temp_sm2 = sm2.CryptSM2(public_key='', private_key=private_key)
    public_key = temp_sm2.public_key
    
    # 保存密钥到文件
    with open(private_key_file, "w") as f:
        f.write(private_key)
    with open(public_key_file, "w") as f:
        f.write(public_key)
    print(f"生成的新私钥已保存到: {private_key_file}")
    print(f"生成的新公钥已保存到: {public_key_file}")
    
    # 创建SM2对象
    sm2_crypt = sm2.CryptSM2(
        public_key=public_key, 
        private_key=private_key
    )
    
    return sm2_crypt

def encrypt_csv(input_data, output_file=None, exclude_columns=None, use_fixed_key=True):
    """加密CSV数据
    
    参数:
        input_data: 可以是CSV文件路径或pandas DataFrame对象
        output_file: 输出文件路径，可选
        exclude_columns: 不加密的列名列表，默认为['stcd', 'tm']
        use_fixed_key: 是否使用固定密钥，默认为True
    
    返回:
        加密后的DataFrame
    """
    if exclude_columns is None:
        exclude_columns = ['stcd', 'tm']
    
    # 设置密钥
    sm2_crypt = setup_keys(use_fixed_key)
    
    # 定义内部加密函数
    def encrypt_value(value, sm2_crypt):
        """加密单个值"""
        try:
            # 检查是否为NaN或None
            if pd.isna(value):
                return "NA_ENCRYPTED"  # 使用特殊标记表示空值
                
            # 将值转换为字符串
            value_str = str(value)
            
            # 对于数值型数据，可以使用一种更简单的加密方式
            if value_str.replace('.', '', 1).isdigit():  # 检查是否为数字（包括小数）
                # 简单加密：对于数值，可以使用一个固定的前缀加上Base64编码
                simple_enc = base64.b64encode(value_str.encode()).decode()
                return f"NUM_ENC_{simple_enc}"
                
            # 使用SM2加密（对于非数值型数据）
            try:
                enc_data = sm2_crypt.encrypt(value_str.encode())
                # 转换为Base64编码便于存储
                return base64.b64encode(enc_data).decode()
            except Exception as e:
                # 如果SM2加密失败，回退到简单加密
                simple_enc = base64.b64encode(value_str.encode()).decode()
                return f"SIMPLE_ENC_{simple_enc}"
        except Exception as e:
            print(f"加密错误: {e}, 值类型: {type(value)}, 值: {value}")
            # 返回一个标记，表示这是未加密的原始值
            return f"ORIGINAL_{value}"
    
    # 读取数据
    if isinstance(input_data, str):
        # 如果是字符串，当作文件路径处理
        df = pd.read_csv(input_data)
    elif isinstance(input_data, pd.DataFrame):
        # 如果是DataFrame，直接使用
        df = input_data.copy()
    else:
        raise TypeError("input_data必须是文件路径或DataFrame对象")
    
    # 对DataFrame中除了排除列以外的列进行加密
    for column in df.columns:
        if column not in exclude_columns:
            print(f"加密列: {column}")
            # 将列转换为字符串类型，确保所有值都可以被加密
            df[column] = df[column].astype(str)
            df[column] = df[column].apply(lambda x: encrypt_value(x, sm2_crypt))
    
    # 如果指定了输出文件，则保存加密后的数据
    if output_file:
        df.to_csv(output_file, index=False)
        print(f"加密后的数据已保存到: {output_file}")
    
    return df

def decrypt_csv(input_data, output_file=None, exclude_columns=None, use_fixed_key=True):
    """解密CSV数据
    
    参数:
        input_data: 可以是CSV文件路径或pandas DataFrame对象
        output_file: 输出文件路径，可选
        exclude_columns: 不解密的列名列表，默认为['stcd', 'tm']
        use_fixed_key: 是否使用固定密钥，默认为True
    
    返回:
        解密后的DataFrame
    """
    if exclude_columns is None:
        exclude_columns = ['stcd', 'tm']
    
    # 设置密钥
    sm2_crypt = setup_keys(use_fixed_key)
    
    # 定义内部解密函数
    def decrypt_value(encrypted_value, sm2_crypt):
        """解密单个值"""
        try:
            # 检查是否为特殊标记
            if encrypted_value == "NA_ENCRYPTED":
                return None
                
            # 检查是否为未加密的原始值
            if isinstance(encrypted_value, str) and encrypted_value.startswith("ORIGINAL_"):
                original_value = encrypted_value[9:]  # 去掉"ORIGINAL_"前缀
                # 尝试转换回数值
                try:
                    if '.' in original_value:
                        return float(original_value)
                    elif original_value.isdigit():
                        return int(original_value)
                    else:
                        return original_value
                except:
                    return original_value
            
            # 检查是否为数值简单加密
            if isinstance(encrypted_value, str) and encrypted_value.startswith("NUM_ENC_"):
                # 解码Base64获取原始数值字符串
                value_str = base64.b64decode(encrypted_value[8:]).decode()
                # 尝试转换回数值
                try:
                    if '.' in value_str:
                        return float(value_str)
                    elif value_str.isdigit():
                        return int(value_str)
                    else:
                        return value_str
                except:
                    return value_str
                    
            # 检查是否为简单加密
            if isinstance(encrypted_value, str) and encrypted_value.startswith("SIMPLE_ENC_"):
                # 解码Base64获取原始字符串
                return base64.b64decode(encrypted_value[11:]).decode()
                
            # 标准SM2解密
            try:
                enc_data = base64.b64decode(encrypted_value)
                dec_data = sm2_crypt.decrypt(enc_data)
                return dec_data.decode()
            except Exception as e:
                print(f"SM2解密失败: {e}, 尝试返回原始值")
                return encrypted_value
        except Exception as e:
            print(f"解密错误: {e}, 值类型: {type(encrypted_value)}, 值: {encrypted_value}")
            return encrypted_value
    
    # 读取数据
    if isinstance(input_data, str):
        # 如果是字符串，当作文件路径处理
        df = pd.read_csv(input_data)
    elif isinstance(input_data, pd.DataFrame):
        # 如果是DataFrame，直接使用
        df = input_data.copy()
    else:
        raise TypeError("input_data必须是文件路径或DataFrame对象")
    
    # 对DataFrame中除了排除列以外的列进行解密
    for column in df.columns:
        if column not in exclude_columns:
            print(f"解密列: {column}")
            df[column] = df[column].apply(lambda x: decrypt_value(x, sm2_crypt))
            
            # 尝试将数值列转换回数值类型
            try:
                # 检查是否可以转换为数值
                numeric_series = pd.to_numeric(df[column], errors='coerce')
                # 如果没有产生太多NaN值，则认为这是数值列
                if numeric_series.isna().sum() / len(numeric_series) < 0.1:  # 少于10%的NaN
                    df[column] = numeric_series
            except:
                pass  # 如果转换失败，保持字符串类型
    
    # 如果指定了输出文件，则保存解密后的数据
    if output_file:
        df.to_csv(output_file, index=False)
        print(f"解密后的数据已保存到: {output_file}")
    
    return df
#
# def main():
#     """主函数，处理命令行参数"""
#     parser = argparse.ArgumentParser(description='CSV文件加密/解密工具')
#     parser.add_argument('action', choices=['encrypt', 'decrypt', 'process'], help='执行的操作: encrypt(加密), decrypt(解密), 或 process(处理两个CSV文件)')
#     parser.add_argument('--input', '-i', help='输入CSV文件路径')
#     parser.add_argument('--output', '-o', help='输出CSV文件路径')
#     parser.add_argument('--exclude', '-e', nargs='+', default=['stcd', 'tm'], help='排除不加密/解密的列名列表')
#     parser.add_argument('--fixed-key', '-f', action='store_true', help='使用固定密钥')
#
#     args = parser.parse_args()
#
#     if args.action == 'encrypt':
#         if not args.input:
#             print("错误: 加密操作需要指定输入文件")
#             return
#         encrypt_csv(args.input, args.output, args.exclude, args.fixed_key)
#     elif args.action == 'decrypt':
#         if not args.input:
#             print("错误: 解密操作需要指定输入文件")
#             return
#         decrypt_csv(args.input, args.output, args.exclude, args.fixed_key)

if __name__ == "__main__":
    # 文件路径
    file1 = "d:\\source pro\\secret\\st_pptn_r_202503051323——时段.csv"
    file2 = "d:\\source pro\\secret\\st_pptn_r_202503251135——整编.csv"
    
    # 为每个文件创建输出路径
    file1_encrypted = "d:\\source pro\\secret\\st_pptn_r_202503051323——时段_加密.csv"
    file1_decrypted = "d:\\source pro\\secret\\st_pptn_r_202503051323——时段_解密.csv"
    file2_encrypted = "d:\\source pro\\secret\\st_pptn_r_202503251135——整编_加密.csv"
    file2_decrypted = "d:\\source pro\\secret\\st_pptn_r_202503251135——整编_解密.csv"
    
    # 处理第一个文件
    print("\n===== 处理文件1: st_pptn_r_202503051323——时段.csv =====")
    # 读取原始数据
    df1 = pd.read_csv(file1)
    print("\n原始数据预览 (前5行):")
    print(df1.head())
    
    # 加密数据
    encrypted_df1 = encrypt_csv(df1, file1_encrypted)
    print("\n加密后数据预览 (前5行):")
    print(encrypted_df1.head())
    
    # 解密数据
    decrypted_df1 = decrypt_csv(encrypted_df1, file1_decrypted)
    print("\n解密后数据预览 (前5行):")
    print(decrypted_df1.head())
    
    # 处理第二个文件
    print("\n\n===== 处理文件2: st_pptn_r_202503251135——整编.csv =====")
    # 读取原始数据
    df2 = pd.read_csv(file2)
    print("\n原始数据预览 (前5行):")
    print(df2.head())
    
    # 加密数据
    encrypted_df2 = encrypt_csv(df2, file2_encrypted)
    print("\n加密后数据预览 (前5行):")
    print(encrypted_df2.head())
    
    # 解密数据
    decrypted_df2 = decrypt_csv(encrypted_df2, file2_decrypted)
    print("\n解密后数据预览 (前5行):")
    print(decrypted_df2.head())
    
    print("\n处理完成！所有文件已保存。")
    print(f"文件1加密结果: {file1_encrypted}")
    print(f"文件1解密结果: {file1_decrypted}")
    print(f"文件2加密结果: {file2_encrypted}")
    print(f"文件2解密结果: {file2_decrypted}")
    
    # 如果需要通过命令行使用，请注释上面的代码，取消下面的注释
    # main()
